import json
import time
import logging
import uuid
import asyncio
from typing import Dict, Any, AsyncGenerator

from fastapi.responses import StreamingResponse, JSONResponse
from playwright.async_api import async_playwright, Page, BrowserContext, Route, TimeoutError as PlaywrightTimeoutError, Error as PlaywrightError
from curl_cffi.requests import AsyncSession

from app.core.config import settings
from app.utils.sse_utils import create_sse_data, create_chat_completion_chunk, DONE_CHUNK

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

class EaseMateProvider:
    UI_MODEL_NAME_MAP = {
        "meta-llama-3.1": "Meta Llama 3.3",
        "claude-3-haiku": "Claude 3 Haiku",
        "gpt-4o-mini": "GPT-4o mini",
        "deepseek-v2": "DeepSeek V3",
        "deepseek-coder-r1": "DeepSeek R1",
        "gemini-1.5-flash": "Gemini 2.0 Flash",
        "gemini-1.5-pro": "Gemini 2.5 Flash",
        "moonshot-v1": "Kimi K2",
        "qwen2-72b-instruct": "Qwen3 235B",
    }

    def __init__(self):
        self.base_url = "https://api.easemate.ai"
        self.exec_operation_url = f"{self.base_url}/api2/stream/exec_operation"
        
        self.playwright = None
        self.browser_context: BrowserContext | None = None
        self.page: Page | None = None
        logger.info("EaseMateProvider 实例已创建。")

    async def initialize_browser(self):
        self.playwright = await async_playwright().start()
        self.browser_context = await self.playwright.chromium.launch_persistent_context(
            user_data_dir="/app/browser_data",
            headless=True,
            args=[
                "--no-sandbox",
                "--disable-setuid-sandbox",
                "--disable-infobars",
                "--disable-blink-features=AutomationControlled",
                "--disable-web-security"
            ],
            user_agent="Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36"
        )
        
        self.page = self.browser_context.pages[0] if self.browser_context.pages else await self.browser_context.new_page()
        
        logger.info("导航到 EaseMate 聊天页面以建立和预热浏览器环境...")
        await self.page.goto("https://www.easemate.ai/webapp/chat", wait_until="networkidle", timeout=60000)
        logger.info("持久化浏览器上下文已启动并导航到目标页面。")

    async def close_browser(self):
        if self.browser_context:
            await self.browser_context.close()
            logger.info("持久化浏览器上下文已关闭。")
        if self.playwright:
            await self.playwright.stop()

    def is_ready(self) -> bool:
        return self.page is not None and not self.page.is_closed()

    async def _handle_login_popup(self):
        """检查并关闭登录弹窗"""
        close_button_selector = "div.ant-modal-wrap > div.ant-modal-content > button.ant-modal-close"
        try:
            close_button = self.page.locator(close_button_selector)
            if await close_button.is_visible(timeout=2000):
                logger.info("检测到登录弹窗，正在尝试关闭...")
                await close_button.click()
                await self.page.wait_for_timeout(500)
                logger.info("登录弹窗已关闭。")
        except PlaywrightTimeoutError:
            logger.info("未检测到登录弹窗。")
        except Exception as e:
            logger.warning(f"关闭登录弹窗时出现问题: {e}")

    async def _stream_generator(self, request_data: Dict[str, Any]) -> AsyncGenerator[bytes, None]:
        request_id = f"chatcmpl-{uuid.uuid4()}"
        model_name = request_data.get("model", "gpt-4o-mini")
        
        await self._handle_login_popup()

        try:
            logger.info(f"用户请求的模型: '{model_name}'")
            current_model_selector = "div.model-select-active span.text-sm"
            await self.page.wait_for_selector(current_model_selector, timeout=15000)
            current_model_on_page = await self.page.inner_text(current_model_selector)
            logger.info(f"页面当前显示的模型: '{current_model_on_page}'")

            target_ui_name = self.UI_MODEL_NAME_MAP.get(model_name)

            if not target_ui_name:
                 logger.warning(f"警告：在UI_MODEL_NAME_MAP中未找到模型 '{model_name}' 的映射。将使用当前页面模型 '{current_model_on_page}'。")
            elif current_model_on_page.lower().strip() != target_ui_name.lower().strip():
                logger.info(f"需要切换模型: 从 '{current_model_on_page}' 切换到 '{target_ui_name}'")
                
                dropdown_selector = "div.model-select-active"
                await self.page.click(dropdown_selector)
                logger.info("已点击模型选择下拉菜单。")
                await self.page.wait_for_timeout(500)

                target_model_selector = f"//div[contains(@class, 'model-item')]//span[text()='{target_ui_name}']"
                await self.page.wait_for_selector(target_model_selector, timeout=5000)
                await self.page.click(target_model_selector)
                logger.info(f"已在列表中点击模型: '{target_ui_name}'。")
                await self.page.wait_for_timeout(500)
            else:
                logger.info("请求的模型已是当前活动模型，无需切换。")

        except PlaywrightTimeoutError as e:
            logger.error(f"错误：在UI中找不到或操作模型时超时。错误: {e}。将使用当前模型继续。")
        except Exception as e:
            logger.error(f"切换模型时发生意外错误: {e}", exc_info=True)
        
        queue = asyncio.Queue()
        interception_future = asyncio.get_running_loop().create_future()

        async def handle_route(route: Route):
            try:
                if "google-analytics.com" in route.request.url or "easeus.recovery-soft.com" in route.request.url:
                    await route.abort()
                    return

                is_target_request = self.exec_operation_url in route.request.url and route.request.method == "POST"

                if is_target_request:
                    if interception_future.done():
                        logger.warning(f"重复的 exec_operation 请求被拦截，忽略。")
                        await route.fulfill(status=200, body="{}")
                        return

                    logger.info(f"成功拦截到浏览器的流式请求: {route.request.url}")
                    try:
                        original_payload = json.loads(route.request.post_data_buffer.decode('utf-8'))
                        logger.info(f"浏览器正在发送的载荷: {json.dumps(original_payload, ensure_ascii=False)}")
                        
                        async with AsyncSession(impersonate="chrome110") as session:
                            async with session.stream(
                                method=route.request.method,
                                url=route.request.url,
                                headers=await route.request.all_headers(),
                                data=route.request.post_data_buffer,
                                timeout=settings.API_REQUEST_TIMEOUT
                            ) as response:
                                response.raise_for_status()
                                logger.info("curl_cffi 成功连接流式接口，准备转发数据...")
                                async for chunk in response.aiter_content():
                                    await queue.put(chunk)
                        
                        if not interception_future.done():
                            interception_future.set_result(True)
                    except Exception as e:
                        logger.error(f"拦截并转发请求时出错: {e}", exc_info=True)
                        if not interception_future.done():
                            interception_future.set_exception(e)
                    finally:
                        await route.fulfill(status=200, body=json.dumps({"code": 200, "message": "Request intercepted and handled by proxy."}))
                else:
                    await route.continue_()
            except PlaywrightError as e:
                if "already handled" not in str(e).lower():
                    logger.warning(f"处理路由时发生Playwright错误: {e}")
            except Exception as e:
                logger.error(f"处理路由时发生未知错误: {e}", exc_info=True)

        try:
            await self.page.route("**/*", handle_route)

            messages = request_data.get("messages", [])
            last_message = messages[-1]['content'] if messages and messages[-1]['role'] == 'user' else "你好"
            
            chat_input_selector = 'textarea[placeholder="Ask me anything…"]'
            send_button_selector = 'button.css-1wchz4a'

            await self.page.wait_for_selector(chat_input_selector, timeout=15000)
            await self.page.fill(chat_input_selector, last_message)
            await self.page.wait_for_timeout(500)
            
            logger.info("点击发送按钮，触发网站自身逻辑...")
            await self.page.click(send_button_selector, force=True)
            
            await asyncio.wait_for(interception_future, timeout=20)

            # --- 核心改进：引入首次和后续超时 ---
            first_chunk_received = False
            FIRST_CHUNK_TIMEOUT = 60.0  # 首次响应的超时时间（秒）
            SUBSEQUENT_CHUNK_TIMEOUT = 3.0 # 后续数据块的超时时间（秒）

            while True:
                try:
                    timeout = SUBSEQUENT_CHUNK_TIMEOUT if first_chunk_received else FIRST_CHUNK_TIMEOUT
                    chunk = await asyncio.wait_for(queue.get(), timeout=timeout)
                    
                    if not first_chunk_received:
                        logger.info("已收到第一个数据块，切换到后续超时策略。")
                        first_chunk_received = True

                    chunk_str = chunk.decode('utf-8', errors='ignore')
                    for line in chunk_str.splitlines():
                        if line.startswith("data:"):
                            data_part = line[len("data:"):].strip()
                            try:
                                data = json.loads(data_part)
                                if data.get("code") == 200:
                                    inner_data_str = data.get("data")
                                    if isinstance(inner_data_str, str):
                                        answer_data = json.loads(inner_data_str)
                                        delta_content = answer_data.get("answer")
                                        if delta_content is not None:
                                            yield create_sse_data(create_chat_completion_chunk(request_id, model_name, delta_content))
                            except Exception:
                                pass
                except asyncio.TimeoutError:
                    if not first_chunk_received:
                        logger.error(f"错误：在 {FIRST_CHUNK_TIMEOUT} 秒内未收到任何数据，模型可能响应超时。")
                    else:
                        logger.info("队列数据消费超时，认为流已结束。")
                    break
            
            yield create_sse_data(create_chat_completion_chunk(request_id, model_name, "", "stop"))
            yield DONE_CHUNK

        except Exception as e:
            logger.error(f"流程中发生严重错误: {e}", exc_info=True)
            error_message = f"内部服务器错误: {str(e)}"
            yield create_sse_data(create_chat_completion_chunk(request_id, model_name, error_message, "stop"))
            yield DONE_CHUNK
        finally:
            await self.page.unroute("**/*", handle_route)
            logger.info("已移除网络拦截器。")

    def chat_completion(self, request_data: Dict[str, Any]):
        return StreamingResponse(self._stream_generator(request_data), media_type="text/event-stream")

    async def get_models(self):
        return JSONResponse(content={
            "object": "list",
            "data": [{"id": name, "object": "model", "created": int(time.time()), "owned_by": "lzA6-prometheus"} for name in settings.MODEL_MAP.keys()]
        })
